// Copyright (C) 2019-2020 Zilliz. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software distributed under the License
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
// or implied. See the License for the specific language governing permissions and limitations under the License.

#include "db/DBImpl.h"
#include "cache/CpuCacheMgr.h"
#include "codecs/Codec.h"
#include "db/IDGenerator.h"
#include "db/SnapshotUtils.h"
#include "db/SnapshotVisitor.h"
#include "db/merge/MergeManagerFactory.h"
#include "db/merge/MergeTask.h"
#include "db/snapshot/CompoundOperations.h"
#include "db/snapshot/ResourceHelper.h"
#include "db/snapshot/ResourceTypes.h"
#include "db/snapshot/Snapshots.h"
#include "insert/MemManagerFactory.h"
#include "knowhere/index/vector_index/helpers/BuilderSuspend.h"
#include "scheduler/Definition.h"
#include "scheduler/SchedInst.h"
#include "scheduler/job/SearchJob.h"
#include "segment/SegmentReader.h"
#include "segment/Utils.h"
#include "utils/Exception.h"
#include "utils/TimeRecorder.h"
#include "value/config/ServerConfig.h"
#include "value/status/ServerStatus.h"

#include <fiu/fiu-local.h>
#include <src/scheduler/job/BuildIndexJob.h>
#include <algorithm>
#include <limits>
#include <unordered_set>
#include <utility>

namespace milvus {
namespace engine {

namespace {
constexpr uint64_t BACKGROUND_METRIC_INTERVAL = 1;
constexpr uint64_t BACKGROUND_INDEX_INTERVAL = 1;
constexpr uint64_t WAIT_BUILD_INDEX_INTERVAL = 5;

static const Status SHUTDOWN_ERROR = Status(DB_ERROR, "Milvus server is shutdown!");
static const Status PERMISSION_ERROR = Status(DB_PERMISSION_ERROR, "Write permission needed!");
}  // namespace

#define CHECK_AVAILABLE        \
    if (!ServiceAvailable()) { \
        return SHUTDOWN_ERROR; \
    }

#define WRITE_PERMISSION_NEEDED_RETURN_STATUS                  \
    if (options_.mode_ == DBOptions::MODE::CLUSTER_READONLY) { \
        return PERMISSION_ERROR;                               \
    }

#define WRITE_PERMISSION_NEEDED_NO_RETURN                      \
    if (options_.mode_ == DBOptions::MODE::CLUSTER_READONLY) { \
        return;                                                \
    }

DBImpl::DBImpl(const DBOptions& options)
    : options_(options), available_(false), merge_thread_pool_(1, 1), index_thread_pool_(1, 1), index_task_tracker_(3) {
    mem_mgr_ = MemManagerFactory::Build(options_);
    merge_mgr_ptr_ = MergeManagerFactory::SSBuild(options_);

    /* watch on storage.auto_flush_interval */
    ConfigMgr::GetInstance().Attach("storage.auto_flush_interval", this);

    DBImpl::Start();
}

DBImpl::~DBImpl() {
    ConfigMgr::GetInstance().Detach("storage.auto_flush_interval", this);

    DBImpl::Stop();
}

bool
DBImpl::ServiceAvailable() {
    return available_.load(std::memory_order_acquire);
}

void
DBImpl::SetAvailable(bool available) {
    available_.store(available, std::memory_order_release);
}

////////////////////////////////////////////////////////////////////////////////
// External APIs
////////////////////////////////////////////////////////////////////////////////
Status
DBImpl::Start() {
    if (ServiceAvailable()) {
        return Status::OK();
    }

    // LOG_ENGINE_TRACE_ << "DB service start";
    SetAvailable(true);

    // server may be closed unexpected, these un-merge files need to be merged when server restart
    // and soft-delete files need to be deleted when server restart
    if (options_.mode_ != DBOptions::MODE::CLUSTER_READONLY) {
        snapshot::IDS_TYPE collection_ids;
        snapshot::Snapshots::GetInstance().GetCollectionIds(collection_ids);
        std::set<int64_t> merge_ids;
        for (auto id : collection_ids) {
            merge_ids.insert(id);
        }
        StartMergeTask(merge_ids, true);
    }

    // for distribute version, some nodes are read only
    if (options_.mode_ != DBOptions::MODE::CLUSTER_READONLY) {
        // background flush thread
        bg_flush_thread_ = std::thread(&DBImpl::TimingFlushThread, this);
    }

    // for distribute version, some nodes are read only
    if (options_.mode_ != DBOptions::MODE::CLUSTER_READONLY) {
        // background build index thread
        bg_index_thread_ = std::thread(&DBImpl::TimingIndexThread, this);
    }
    bg_metric_thread_ = std::thread(&DBImpl::BackgroundMetricThread, this);

    return Status::OK();
}

Status
DBImpl::Stop() {
    if (!ServiceAvailable()) {
        return Status::OK();
    }

    SetAvailable(false);

    if (options_.mode_ != DBOptions::MODE::CLUSTER_READONLY) {
        // flush all without merge
        InternalFlush("", false);

        // wait flush thread finish
        swn_flush_.Notify();
        bg_flush_thread_.join();
        LOG_ENGINE_DEBUG_ << "DBImpl::Stop bg_flush_thread_.join()";

        WaitMergeFileFinish();
        LOG_ENGINE_DEBUG_ << "DBImpl::Stop WaitMergeFileFinish";
        swn_index_.Notify();
        index_req_swn_.Notify();
        bg_index_thread_.join();
        LOG_ENGINE_DEBUG_ << "DBImpl::Stop bg_index_thread_.join()";
        swn_metric_.Notify();
        bg_metric_thread_.join();
    } else {
        swn_metric_.Notify();
        bg_metric_thread_.join();
    }

    // LOG_ENGINE_TRACE_ << "DB service stop";
    return Status::OK();
}

Status
DBImpl::CreateCollection(const snapshot::CreateCollectionContext& context) {
    WRITE_PERMISSION_NEEDED_RETURN_STATUS;
    CHECK_AVAILABLE

    auto ctx = context;

    // default id is auto-generated
    auto params = ctx.collection->GetParams();
    if (params.find(PARAM_UID_AUTOGEN) == params.end()) {
        params[PARAM_UID_AUTOGEN] = true;
    }
    ctx.collection->SetParams(params);

    // check uid existence
    snapshot::FieldPtr uid_field;
    for (auto& pair : ctx.fields_schema) {
        if (pair.first->GetName() == FIELD_UID) {
            uid_field = pair.first;
            break;
        }
    }

    // add uid field if not specified
    if (uid_field == nullptr) {
        uid_field = std::make_shared<snapshot::Field>(FIELD_UID, 0, DataType::INT64);
    }

    // define uid elements
    auto bloom_filter_element = std::make_shared<snapshot::FieldElement>(
        0, 0, ELEMENT_BLOOM_FILTER, milvus::engine::FieldElementType::FET_BLOOM_FILTER);
    auto delete_doc_element = std::make_shared<snapshot::FieldElement>(
        0, 0, ELEMENT_DELETED_DOCS, milvus::engine::FieldElementType::FET_DELETED_DOCS);
    ctx.fields_schema[uid_field] = {bloom_filter_element, delete_doc_element};

    auto op = std::make_shared<snapshot::CreateCollectionOperation>(ctx);
    return op->Push();
}

Status
DBImpl::DropCollection(const std::string& collection_name) {
    WRITE_PERMISSION_NEEDED_RETURN_STATUS;
    CHECK_AVAILABLE

    LOG_ENGINE_DEBUG_ << "Prepare to drop collection " << collection_name;

    snapshot::ScopedSnapshotT ss;
    auto& snapshots = snapshot::Snapshots::GetInstance();
    STATUS_CHECK(snapshots.GetSnapshot(ss, collection_name));

    // erase insert buffer of this collection
    mem_mgr_->EraseMem(ss->GetCollectionId());

    // erase cache
    ClearCollectionCache(ss, options_.meta_.path_);

    // clear index failed retry map of this collection
    index_task_tracker_.ClearFailedRecords(collection_name);

    return snapshots.DropCollection(ss->GetCollectionId(), std::numeric_limits<snapshot::LSN_TYPE>::max());
}

Status
DBImpl::HasCollection(const std::string& collection_name, bool& has_or_not) {
    CHECK_AVAILABLE

    snapshot::ScopedSnapshotT ss;
    auto status = snapshot::Snapshots::GetInstance().GetSnapshot(ss, collection_name);
    has_or_not = status.ok();

    return Status::OK();
}

Status
DBImpl::ListCollections(std::vector<std::string>& names) {
    CHECK_AVAILABLE

    names.clear();
    return snapshot::Snapshots::GetInstance().GetCollectionNames(names);
}

Status
DBImpl::GetCollectionInfo(const std::string& collection_name, snapshot::CollectionPtr& collection,
                          snapshot::FieldElementMappings& fields_schema) {
    CHECK_AVAILABLE

    snapshot::ScopedSnapshotT ss;
    STATUS_CHECK(snapshot::Snapshots::GetInstance().GetSnapshot(ss, collection_name));

    collection = ss->GetCollection();
    auto& fields = ss->GetResources<snapshot::Field>();
    for (auto& kv : fields) {
        fields_schema[kv.second.Get()] = ss->GetFieldElementsByField(kv.second->GetName());
    }
    return Status::OK();
}

Status
DBImpl::GetCollectionStats(const std::string& collection_name, milvus::json& collection_stats) {
    CHECK_AVAILABLE

    STATUS_CHECK(GetSnapshotInfo(collection_name, collection_stats));
    return Status::OK();
}

Status
DBImpl::CountEntities(const std::string& collection_name, int64_t& row_count) {
    CHECK_AVAILABLE

    snapshot::ScopedSnapshotT ss;
    STATUS_CHECK(snapshot::Snapshots::GetInstance().GetSnapshot(ss, collection_name));

    row_count = ss->GetCollectionCommit()->GetRowCount();
    return Status::OK();
}

Status
DBImpl::CreatePartition(const std::string& collection_name, const std::string& partition_name) {
    WRITE_PERMISSION_NEEDED_RETURN_STATUS;
    CHECK_AVAILABLE

    snapshot::ScopedSnapshotT ss;
    STATUS_CHECK(snapshot::Snapshots::GetInstance().GetSnapshot(ss, collection_name));

    snapshot::LSN_TYPE lsn = 0;
    snapshot::OperationContext context;
    context.lsn = lsn;
    auto op = std::make_shared<snapshot::CreatePartitionOperation>(context, ss);

    snapshot::PartitionContext p_ctx;
    p_ctx.name = partition_name;
    snapshot::PartitionPtr partition;
    STATUS_CHECK(op->CommitNewPartition(p_ctx, partition));
    return op->Push();
}

Status
DBImpl::DropPartition(const std::string& collection_name, const std::string& partition_name) {
    WRITE_PERMISSION_NEEDED_RETURN_STATUS;
    CHECK_AVAILABLE

    snapshot::ScopedSnapshotT ss;
    STATUS_CHECK(snapshot::Snapshots::GetInstance().GetSnapshot(ss, collection_name));

    auto partition = ss->GetPartition(partition_name);
    if (partition != nullptr) {
        // erase insert buffer of this partition
        mem_mgr_->EraseMem(ss->GetCollectionId(), partition->GetID());

        // erase cache
        ClearPartitionCache(ss, options_.meta_.path_, partition->GetID());
    }

    snapshot::PartitionContext context;
    context.name = partition_name;
    auto op = std::make_shared<snapshot::DropPartitionOperation>(context, ss);
    return op->Push();
}

Status
DBImpl::HasPartition(const std::string& collection_name, const std::string& partition_tag, bool& exist) {
    CHECK_AVAILABLE

    snapshot::ScopedSnapshotT ss;
    STATUS_CHECK(snapshot::Snapshots::GetInstance().GetSnapshot(ss, collection_name));

    auto partition_tags = std::move(ss->GetPartitionNames());
    for (auto& tag : partition_tags) {
        if (tag == partition_tag) {
            exist = true;
            return Status::OK();
        }
    }

    exist = false;
    return Status::OK();
}

Status
DBImpl::ListPartitions(const std::string& collection_name, std::vector<std::string>& partition_names) {
    CHECK_AVAILABLE

    snapshot::ScopedSnapshotT ss;
    STATUS_CHECK(snapshot::Snapshots::GetInstance().GetSnapshot(ss, collection_name));

    partition_names = std::move(ss->GetPartitionNames());
    return Status::OK();
}

Status
DBImpl::CreateIndex(const std::shared_ptr<server::Context>& context, const std::string& collection_name,
                    const std::string& field_name, const CollectionIndex& index) {
    WRITE_PERMISSION_NEEDED_RETURN_STATUS;
    CHECK_AVAILABLE
    SetThreadName("create_index");
    LOG_ENGINE_DEBUG_ << "Create index for collection: " << collection_name << " field: " << field_name;

    // step 1: wait merge file thread finished to avoid duplicate data bug
    auto status = Flush();
    WaitMergeFileFinish();  // let merge file thread finish

    // step 2: compare old index and new index, drop old index, set new index
    CollectionIndex new_index = index;
    CollectionIndex old_index;
    STATUS_CHECK(GetSnapshotIndex(collection_name, field_name, old_index));

    if (!utils::IsSameIndex(old_index, new_index)) {
        DropIndex(collection_name, field_name);

        WaitMergeFileFinish();  // let merge file thread finish since DropIndex start a merge task

        // create field element for new index
        status = SetSnapshotIndex(collection_name, field_name, new_index);
        if (!status.ok()) {
            return status;
        }
    }

    if (engine::utils::IsFlatIndexType(index.index_type_)) {
        return Status::OK();  // for IDMAP type, no need to create index
    }

    // step 3: merge segments before create index, since there could be some small segments just flushed
    {
        snapshot::ScopedSnapshotT latest_ss;
        STATUS_CHECK(snapshot::Snapshots::GetInstance().GetSnapshot(latest_ss, collection_name));
        std::set<int64_t> collection_ids = {latest_ss->GetCollectionId()};
        StartMergeTask(collection_ids, true);  // start force-merge task
        WaitMergeFileFinish();                 // let force-merge file thread finish
    }

    // clear index failed retry map of this collection
    index_task_tracker_.ClearFailedRecords(collection_name);

    // step 4: iterate segments need to be build index, wait until all segments are built
    while (true) {
        // server is going to shutdown, quit this thread
        if (!ServiceAvailable()) {
            LOG_ENGINE_DEBUG_ << "Build index stopped since DB service going to exit";
            break;
        }

        // start background build index thread
        std::vector<std::string> collection_names = {collection_name};
        StartBuildIndexTask(collection_names, true);

        // check if all segments are built
        SnapshotVisitor ss_visitor(collection_name);
        snapshot::IDS_TYPE segment_ids;
        ss_visitor.SegmentsToIndex(field_name, segment_ids, true);
        if (segment_ids.empty()) {
            break;  // all segments build index finished
        }

        index_task_tracker_.IgnoreFailedSegments(collection_name, segment_ids);
        if (segment_ids.empty()) {
            break;  // some segments failed to build index, and ignored
        }

        index_req_swn_.Wait_For(std::chrono::seconds(1));

        // client break the connection, no need to block, check every 1 second
        if (context && context->IsConnectionBroken()) {
            LOG_ENGINE_DEBUG_ << "Client connection broken, build index in background";
            break;  // just break, not return, continue to update partitions files to to_index
        }
    }

    // step 5: return error message to client if any segment failed to build index
    SegmentFailedMap failed_map;
    index_task_tracker_.GetFailedRecords(collection_name, failed_map);
    if (!failed_map.empty()) {
        auto pair = failed_map.begin();
        std::string msg =
            "Failed to build segment " + std::to_string(pair->first) + " for collection " + collection_name;
        msg += ", reason: ";
        msg += pair->second.message();
        LOG_ENGINE_ERROR_ << msg;

        return Status(DB_ERROR, msg);
    }

    return Status::OK();
}

Status
DBImpl::DropIndex(const std::string& collection_name, const std::string& field_name) {
    WRITE_PERMISSION_NEEDED_RETURN_STATUS;
    CHECK_AVAILABLE

    LOG_ENGINE_DEBUG_ << "Drop index for collection: " << collection_name << " field: " << field_name;

    STATUS_CHECK(DeleteSnapshotIndex(collection_name, field_name));

    snapshot::ScopedSnapshotT ss;
    STATUS_CHECK(snapshot::Snapshots::GetInstance().GetSnapshot(ss, collection_name));

    ClearIndexCache(ss, options_.meta_.path_, field_name);

    std::set<int64_t> collection_ids = {ss->GetCollectionId()};
    StartMergeTask(collection_ids, true);

    return Status::OK();
}

Status
DBImpl::DescribeIndex(const std::string& collection_name, const std::string& field_name, CollectionIndex& index) {
    CHECK_AVAILABLE

    LOG_ENGINE_DEBUG_ << "Describe index for collection: " << collection_name << " field: " << field_name;

    STATUS_CHECK(GetSnapshotIndex(collection_name, field_name, index));

    return Status::OK();
}

Status
DBImpl::Insert(const std::string& collection_name, const std::string& partition_name, DataChunkPtr& data_chunk,
               idx_t op_id) {
    WRITE_PERMISSION_NEEDED_RETURN_STATUS;
    CHECK_AVAILABLE
    ScopedTimer scope_timer([data_chunk, this](double latency) {
        auto size = utils::GetSizeOfChunk(data_chunk);
        this->insert_entities_size_gauge_.Set(size / latency);
    });
    if (data_chunk == nullptr) {
        return Status(DB_ERROR, "Null pointer");
    }

    snapshot::ScopedSnapshotT ss;
    STATUS_CHECK(snapshot::Snapshots::GetInstance().GetSnapshot(ss, collection_name));

    auto partition = ss->GetPartition(partition_name);
    if (partition == nullptr) {
        return Status(DB_NOT_FOUND, "Fail to get partition " + partition_name);
    }

    auto id_field = ss->GetField(FIELD_UID);
    if (id_field == nullptr) {
        return Status(DB_ERROR, "Field '_id' not found");
    }

    // check field names
    auto field_names = ss->GetFieldNames();
    std::unordered_set<std::string> collection_field_names;
    for (auto& name : field_names) {
        collection_field_names.insert(name);
    }
    collection_field_names.erase(engine::FIELD_UID);

    std::unordered_set<std::string> chunk_field_names;
    for (auto& pair : data_chunk->fixed_fields_) {
        chunk_field_names.insert(pair.first);
    }
    for (auto& pair : data_chunk->variable_fields_) {
        chunk_field_names.insert(pair.first);
    }
    chunk_field_names.erase(engine::FIELD_UID);

    if (collection_field_names.size() != chunk_field_names.size()) {
        std::string msg = "Collection has " + std::to_string(collection_field_names.size()) +
                          " fields while the insert data has " + std::to_string(chunk_field_names.size()) + " fields";
        return Status(DB_ERROR, msg);
    } else {
        for (auto& name : chunk_field_names) {
            if (collection_field_names.find(name) == collection_field_names.end()) {
                std::string msg = "The field " + name + " is not defined in collection mapping";
                return Status(DB_ERROR, msg);
            }
        }
    }

    // check id field existence
    auto& params = ss->GetCollection()->GetParams();
    bool auto_genid = true;
    if (params.find(PARAM_UID_AUTOGEN) != params.end()) {
        auto_genid = params[PARAM_UID_AUTOGEN];
    }

    FIXEDX_FIELD_MAP& fields = data_chunk->fixed_fields_;
    auto pair = fields.find(engine::FIELD_UID);
    if (auto_genid) {
        // id is auto generated, but client provides id, return error
        if (pair != fields.end() && pair->second != nullptr) {
            return Status(DB_ERROR, "Field '_id' is auto increment, no need to provide id");
        }
    } else {
        // id is not auto generated, but client doesn't provide id, return error
        if (pair == fields.end() || pair->second == nullptr) {
            return Status(DB_ERROR, "Field '_id' is user defined");
        }
    }

    // consume the data chunk
    DataChunkPtr consume_chunk = std::make_shared<DataChunk>();
    consume_chunk->count_ = data_chunk->count_;
    consume_chunk->fixed_fields_.swap(data_chunk->fixed_fields_);
    consume_chunk->variable_fields_.swap(data_chunk->variable_fields_);

    // generate id
    if (auto_genid) {
        LOG_SERVER_DEBUG_ << "Auto generate entities id";

        SafeIDGenerator& id_generator = SafeIDGenerator::GetInstance();
        IDNumbers ids;
        STATUS_CHECK(id_generator.GetNextIDNumbers(consume_chunk->count_, ids));
        BinaryDataPtr id_data = std::make_shared<BinaryData>();
        id_data->data_.resize(ids.size() * sizeof(int64_t));
        memcpy(id_data->data_.data(), ids.data(), ids.size() * sizeof(int64_t));
        consume_chunk->fixed_fields_[engine::FIELD_UID] = id_data;
        data_chunk->fixed_fields_[engine::FIELD_UID] = id_data;  // return generated id to customer;
    } else {
        BinaryDataPtr id_data = std::make_shared<BinaryData>();
        id_data->data_ = consume_chunk->fixed_fields_[engine::FIELD_UID]->data_;
        data_chunk->fixed_fields_[engine::FIELD_UID] = id_data;  // return the id created by client
    }

    // do insert
    int64_t segment_row_limit = 0;
    GetSegmentRowLimit(ss->GetCollection(), segment_row_limit);

    int64_t collection_id = ss->GetCollectionId();
    int64_t partition_id = partition->GetID();

    std::vector<DataChunkPtr> chunks;
    STATUS_CHECK(utils::SplitChunk(consume_chunk, segment_row_limit, chunks));

    LOG_ENGINE_DEBUG_ << "Insert entities into mem manager";
    for (auto& chunk : chunks) {
        auto status = mem_mgr_->InsertEntities(collection_id, partition_id, chunk, op_id);
        if (!status.ok()) {
            return status;
        }

        std::set<int64_t> collection_ids;
        if (mem_mgr_->RequireFlush(collection_ids)) {
            LOG_ENGINE_DEBUG_ << "Insert buffer size exceeds limit. Force flush";
            InternalFlush();
        }
    }

    return Status::OK();
}

Status
DBImpl::GetEntityByID(const std::string& collection_name, const IDNumbers& id_array,
                      const std::vector<std::string>& field_names, std::vector<bool>& valid_row,
                      DataChunkPtr& data_chunk) {
    CHECK_AVAILABLE

    snapshot::ScopedSnapshotT ss;
    STATUS_CHECK(snapshot::Snapshots::GetInstance().GetSnapshot(ss, collection_name));

    std::string dir_root = options_.meta_.path_;
    valid_row.resize(id_array.size(), false);
    auto handler =
        std::make_shared<GetEntityByIdSegmentHandler>(nullptr, ss, dir_root, id_array, field_names, valid_row);
    handler->Iterate();
    STATUS_CHECK(handler->GetStatus());

    data_chunk = handler->data_chunk_;
    return Status::OK();
}

Status
DBImpl::DeleteEntityByID(const std::string& collection_name, const engine::IDNumbers& entity_ids, idx_t op_id) {
    WRITE_PERMISSION_NEEDED_RETURN_STATUS;
    CHECK_AVAILABLE

    snapshot::ScopedSnapshotT ss;
    auto status = snapshot::Snapshots::GetInstance().GetSnapshot(ss, collection_name);
    if (!status.ok()) {
        LOG_ENGINE_DEBUG_ << "Failed to get snapshot: " << status.message();
        return status;
    }

    status = mem_mgr_->DeleteEntities(ss->GetCollectionId(), entity_ids, op_id);
    if (!status.ok()) {
        return status;
    }

    std::set<int64_t> collection_ids;
    if (mem_mgr_->RequireFlush(collection_ids)) {
        if (collection_ids.find(ss->GetCollectionId()) != collection_ids.end()) {
            LOG_ENGINE_DEBUG_ << "Delete count in buffer exceeds limit. Force flush";
            InternalFlush(collection_name);
        }
    }

    return Status::OK();
}

Status
DBImpl::Query(const server::ContextPtr& context, const query::QueryPtr& query_ptr, engine::QueryResultPtr& result) {
    CHECK_AVAILABLE
    SetThreadName("query");
    TimeRecorderAuto rc("DBImpl::Query");

    if (!query_ptr->root) {
        return Status{DB_ERROR, "BinaryQuery is null"};
    }

    auto vector_param = query_ptr->vectors.begin()->second;
    auto nq = vector_param->nq;
    ScopedTimer scoped_timer([nq, this](double latency) {
        for (int64_t i = 0; i < nq; ++i) {
            this->query_count_summary_.Observe(latency / nq);
            this->query_response_summary_.Observe(latency);
        }
    });
    snapshot::ScopedSnapshotT ss;
    STATUS_CHECK(snapshot::Snapshots::GetInstance().GetSnapshot(ss, query_ptr->collection_id));

    SnapshotVisitor ss_visitor(ss);
    snapshot::IDS_TYPE segment_ids;
    STATUS_CHECK(ss_visitor.SegmentsToSearch(query_ptr->partitions, segment_ids));
    rc.RecordSection("segments to search: " + std::to_string(segment_ids.size()));

    scheduler::SearchJobPtr job = std::make_shared<scheduler::SearchJob>(nullptr, ss, options_, query_ptr, segment_ids);

    cache::CpuCacheMgr::GetInstance().PrintInfo();  // print cache info before query

    // SuspendIfFirst();

    /* put search job to scheduler and wait job finish */
    scheduler::JobMgrInst::GetInstance()->Put(job);
    job->WaitFinish();

    // ResumeIfLast();

    cache::CpuCacheMgr::GetInstance().PrintInfo();  // print cache info after query

    if (!job->status().ok()) {
        return job->status();
    }

    if (job->query_result()) {
        result = job->query_result();
    }
    rc.RecordSection("execute query");

    // step 4: get entities by result ids
    std::vector<bool> valid_row;
    if (!query_ptr->field_names.empty()) {
        STATUS_CHECK(GetEntityByID(query_ptr->collection_id, result->result_ids_, query_ptr->field_names, valid_row,
                                   result->data_chunk_));
        rc.RecordSection("get entities");
    }

    return Status::OK();
}

Status
DBImpl::ListIDInSegment(const std::string& collection_name, int64_t segment_id, IDNumbers& entity_ids) {
    CHECK_AVAILABLE

    snapshot::ScopedSnapshotT ss;
    STATUS_CHECK(snapshot::Snapshots::GetInstance().GetSnapshot(ss, collection_name));

    auto read_visitor = engine::SegmentVisitor::Build(ss, segment_id);
    if (!read_visitor) {
        return Status(SERVER_FILE_NOT_FOUND, "Segment not exist");
    }
    segment::SegmentReaderPtr segment_reader =
        std::make_shared<segment::SegmentReader>(options_.meta_.path_, read_visitor);

    STATUS_CHECK(segment_reader->LoadUids(entity_ids));

    // remove delete id from the id list
    segment::DeletedDocsPtr deleted_docs_ptr;
    segment_reader->LoadDeletedDocs(deleted_docs_ptr);
    if (deleted_docs_ptr) {
        // sorted-merge entities id and deleted offsets
        const std::vector<offset_t>& delete_offsets = deleted_docs_ptr->GetDeletedDocs();
        IDNumbers result_ids;
        segment::GetIDWithoutDeleted(entity_ids, delete_offsets, result_ids);
        entity_ids.swap(result_ids);
    }

    return Status::OK();
}

Status
DBImpl::LoadCollection(const server::ContextPtr& context, const std::string& collection_name,
                       const std::vector<std::string>& field_names, bool force) {
    CHECK_AVAILABLE

    snapshot::ScopedSnapshotT ss;
    STATUS_CHECK(snapshot::Snapshots::GetInstance().GetSnapshot(ss, collection_name));

    auto handler = std::make_shared<LoadCollectionHandler>(nullptr, ss, options_.meta_.path_, field_names, force);
    handler->Iterate();
    STATUS_CHECK(handler->GetStatus());

    return Status::OK();
}

Status
DBImpl::Flush(const std::string& collection_name) {
    WRITE_PERMISSION_NEEDED_RETURN_STATUS;
    CHECK_AVAILABLE

    Status status;
    bool has_collection = false;
    status = HasCollection(collection_name, has_collection);
    if (!status.ok()) {
        return status;
    }
    if (!has_collection) {
        LOG_ENGINE_ERROR_ << "Collection to flush does not exist: " << collection_name;
        return Status(DB_NOT_FOUND, "Collection to flush does not exist");
    }

    LOG_ENGINE_DEBUG_ << "Begin flush collection: " << collection_name;
    InternalFlush(collection_name);
    LOG_ENGINE_DEBUG_ << "End flush collection: " << collection_name;

    return status;
}

Status
DBImpl::Flush() {
    WRITE_PERMISSION_NEEDED_RETURN_STATUS;
    CHECK_AVAILABLE

    LOG_ENGINE_DEBUG_ << "Begin flush all collections";
    InternalFlush();
    LOG_ENGINE_DEBUG_ << "End flush all collections";

    return Status::OK();
}

Status
DBImpl::Compact(const std::shared_ptr<server::Context>& context, const std::string& collection_name, double threshold) {
    WRITE_PERMISSION_NEEDED_RETURN_STATUS;
    CHECK_AVAILABLE

    LOG_ENGINE_DEBUG_ << "Before compacting, wait for build index thread to finish...";
    const std::lock_guard<std::mutex> index_lock(build_index_mutex_);
    const std::lock_guard<std::mutex> merge_lock(flush_merge_compact_mutex_);

    Status status;
    bool has_collection = false;
    status = HasCollection(collection_name, has_collection);
    if (!status.ok()) {
        return status;
    }
    if (!has_collection) {
        LOG_ENGINE_ERROR_ << "Collection to compact does not exist: " << collection_name;
        return Status(DB_NOT_FOUND, "Collection to compact does not exist");
    }

    snapshot::ScopedSnapshotT latest_ss;
    status = snapshot::Snapshots::GetInstance().GetSnapshot(latest_ss, collection_name);
    if (!status.ok()) {
        return status;
    }

    auto& segments = latest_ss->GetResources<snapshot::Segment>();
    for (auto& kv : segments) {
        // client break the connection, no need to continue
        if (context && context->IsConnectionBroken()) {
            LOG_ENGINE_DEBUG_ << "Client connection broken, stop compact operation";
            break;
        }

        snapshot::ID_TYPE segment_id = kv.first;
        auto read_visitor = engine::SegmentVisitor::Build(latest_ss, segment_id);
        segment::SegmentReaderPtr segment_reader =
            std::make_shared<segment::SegmentReader>(options_.meta_.path_, read_visitor);

        segment::DeletedDocsPtr deleted_docs;
        segment_reader->LoadDeletedDocs(deleted_docs);
        if (deleted_docs == nullptr) {
            continue;  // no deleted docs, no need to compact
        }

        // the segment row count is zero, drop it
        auto segment_commit = latest_ss->GetSegmentCommitBySegmentId(segment_id);
        auto row_count = segment_commit->GetRowCount();
        if (row_count == 0) {
            snapshot::OperationContext drop_seg_context;
            auto seg = latest_ss->GetResource<snapshot::Segment>(segment_id);
            drop_seg_context.prev_segment = seg;
            auto drop_op = std::make_shared<snapshot::DropSegmentOperation>(drop_seg_context, latest_ss);
            status = drop_op->Push();
            if (!status.ok()) {
                LOG_ENGINE_ERROR_ << "Compact failed for segment " << segment_reader->GetSegmentPath() << ": "
                                  << status.message();
            }
            continue;
        }

        // delete rate less than threshold, skip compact
        auto deleted_count = (double)(deleted_docs->GetCount());
        if (deleted_count / (row_count + deleted_count) < threshold) {
            continue;  // no need to compact
        }

        // compact segment, the compact action is same as merge
        snapshot::IDS_TYPE ids = {segment_id};
        MergeTask merge_task(options_, latest_ss, ids);
        status = merge_task.Execute();
        if (!status.ok()) {
            LOG_ENGINE_ERROR_ << "Compact failed for segment " << segment_reader->GetSegmentPath() << ": "
                              << status.message();
            continue;  // skip this file and try compact next one
        }
    }

    return status;
}

////////////////////////////////////////////////////////////////////////////////
// Internal APIs
////////////////////////////////////////////////////////////////////////////////
void
DBImpl::InternalFlush(const std::string& collection_name, bool merge) {
    Status status;
    std::set<int64_t> flushed_collection_ids;
    if (!collection_name.empty()) {
        // flush one collection
        snapshot::ScopedSnapshotT ss;
        status = snapshot::Snapshots::GetInstance().GetSnapshot(ss, collection_name);
        if (!status.ok()) {
            LOG_WAL_ERROR_ << "Failed to get snapshot: " << status.message();
            return;
        }

        {
            const std::lock_guard<std::mutex> lock(flush_merge_compact_mutex_);
            int64_t collection_id = ss->GetCollectionId();
            status = mem_mgr_->Flush(collection_id);
            if (!status.ok()) {
                return;
            }
            flushed_collection_ids.insert(collection_id);
        }
    } else {
        // flush all collections
        {
            const std::lock_guard<std::mutex> lock(flush_merge_compact_mutex_);
            status = mem_mgr_->Flush(flushed_collection_ids);
            if (!status.ok()) {
                return;
            }
        }
    }

    if (merge) {
        StartMergeTask(flushed_collection_ids, false);
    }
}

void
DBImpl::TimingFlushThread() {
    SetThreadName("timing_flush");
    while (true) {
        if (!ServiceAvailable()) {
            LOG_ENGINE_DEBUG_ << "DB background flush thread exit";
            break;
        }

        InternalFlush();
        if (options_.auto_flush_interval_ > 0) {
            swn_flush_.Wait_For(std::chrono::seconds(options_.auto_flush_interval_));
        } else {
            swn_flush_.Wait();
        }
    }
}

void
DBImpl::StartBuildIndexTask(const std::vector<std::string>& collection_names, bool force_build) {
    if (collection_names.empty()) {
        return;  // no need to start thread
    }

    // build index has been finished?
    {
        std::lock_guard<std::mutex> lck(index_result_mutex_);
        if (!index_thread_results_.empty()) {
            std::chrono::milliseconds span(10);
            if (index_thread_results_.back().wait_for(span) == std::future_status::ready) {
                index_thread_results_.pop_back();
            }
        }
    }

    // add new build index task
    {
        std::lock_guard<std::mutex> lck(index_result_mutex_);
        if (index_thread_results_.empty()) {
            index_thread_results_.push_back(
                index_thread_pool_.enqueue(&DBImpl::BackgroundBuildIndexTask, this, collection_names, force_build));
        }
    }
}

void
DBImpl::BackgroundBuildIndexTask(std::vector<std::string> collection_names, bool force_build) {
    SetThreadName("build_index");

    std::unique_lock<std::mutex> lock(build_index_mutex_);

    for (const auto& collection_name : collection_names) {
        snapshot::ScopedSnapshotT latest_ss;
        auto status = snapshot::Snapshots::GetInstance().GetSnapshot(latest_ss, collection_name);
        if (!status.ok()) {
            return;
        }
        SnapshotVisitor ss_visitor(latest_ss);

        snapshot::IDS_TYPE segment_ids;
        ss_visitor.SegmentsToIndex("", segment_ids, force_build);
        if (segment_ids.empty()) {
            continue;
        }

        // check index retry times
        index_task_tracker_.IgnoreFailedSegments(collection_name, segment_ids);
        if (segment_ids.empty()) {
            continue;
        }

        // start build index job
        // build one segment for each time, for two reasons:
        // 1. we don't need to wait all segments index finish when milvus server stop
        // 2. avoid build index for deleted segments
        snapshot::IDS_TYPE segment_to_build = {segment_ids[0]};
        LOG_ENGINE_DEBUG_ << "Create BuildIndexJob for segment " << segment_to_build[0] << " of " << collection_name;
        cache::CpuCacheMgr::GetInstance().PrintInfo();  // print cache info before build index
        scheduler::BuildIndexJobPtr job =
            std::make_shared<scheduler::BuildIndexJob>(latest_ss, options_, segment_to_build);

        IncreaseLiveBuildTaskNum();
        scheduler::JobMgrInst::GetInstance()->Put(job);
        job->WaitFinish();
        DecreaseLiveBuildTaskNum();

        cache::CpuCacheMgr::GetInstance().PrintInfo();  // print cache info after build index

        // record failed segments, avoid build index hang
        auto failed_segments = job->GetFailedSegments();
        index_task_tracker_.MarkFailedSegments(collection_name, failed_segments);

        if (!job->status().ok()) {
            LOG_ENGINE_ERROR_ << job->status().message();
        }

        // notify index request to return (if all segments index has been done)
        index_req_swn_.Notify();

        // quit this thread if the milvus server is going to shutdown
        if (!ServiceAvailable()) {
            LOG_ENGINE_DEBUG_ << "DB background build index thread exit";
            break;
        }
    }
}

void
DBImpl::TimingIndexThread() {
    SetThreadName("timing_index");
    while (true) {
        if (!ServiceAvailable()) {
            WaitMergeFileFinish();
            WaitBuildIndexFinish();

            LOG_ENGINE_DEBUG_ << "DB background timing index thread exit";
            break;
        }

        swn_index_.Wait_For(std::chrono::seconds(BACKGROUND_INDEX_INTERVAL));

        std::vector<std::string> collection_names;
        snapshot::Snapshots::GetInstance().GetCollectionNames(collection_names);
        WaitMergeFileFinish();
        StartBuildIndexTask(collection_names, false);
    }
}

void
DBImpl::WaitBuildIndexFinish() {
    //    LOG_ENGINE_DEBUG_ << "Begin WaitBuildIndexFinish";
    std::lock_guard<std::mutex> lck(index_result_mutex_);
    for (auto& iter : index_thread_results_) {
        iter.wait();
    }
    //    LOG_ENGINE_DEBUG_ << "End WaitBuildIndexFinish";
}

void
DBImpl::StartMergeTask(const std::set<int64_t>& collection_ids, bool force_merge_all) {
    if (collection_ids.empty()) {
        return;  // no need to start thread
    }

    // merge task has been finished?
    {
        std::lock_guard<std::mutex> lck(merge_result_mutex_);
        if (!merge_thread_results_.empty()) {
            std::chrono::milliseconds span(10);
            if (merge_thread_results_.back().wait_for(span) == std::future_status::ready) {
                merge_thread_results_.pop_back();
            }
        }
    }

    // add new merge task
    {
        std::lock_guard<std::mutex> lck(merge_result_mutex_);
        if (merge_thread_results_.empty()) {
            // start merge file thread
            merge_thread_results_.push_back(
                merge_thread_pool_.enqueue(&DBImpl::BackgroundMerge, this, collection_ids, force_merge_all));
        }
    }
}

void
DBImpl::BackgroundMerge(std::set<int64_t> collection_ids, bool force_merge_all) {
    SetThreadName("merge");

    for (auto& collection_id : collection_ids) {
        const std::lock_guard<std::mutex> lock(flush_merge_compact_mutex_);

        MergeStrategyType type = force_merge_all ? MergeStrategyType::ADAPTIVE : MergeStrategyType::LAYERED;
        auto status = merge_mgr_ptr_->MergeSegments(collection_id, type);
        if (!status.ok()) {
            LOG_ENGINE_ERROR_ << "Failed to get merge files for collection id: " << collection_id
                              << " reason:" << status.message();
        }

        if (!ServiceAvailable()) {
            LOG_ENGINE_DEBUG_ << "Server will shutdown, skip merge action for collection id: " << collection_id;
            break;
        }
    }
}

void
DBImpl::BackgroundMetricThread() {
    SetThreadName("metric_thread");
    while (true) {
        if (!ServiceAvailable()) {
            LOG_ENGINE_DEBUG_ << "DB background metric thread exit";
            break;
        }
        size_t data_size = 0;
        auto status = GetDataSize(data_size);
        data_size_gauge_.Set(data_size);
        if (!status.ok()) {
            LOG_ENGINE_ERROR_ << "Server get data size failed";
        }

        swn_metric_.Wait_For(std::chrono::seconds(BACKGROUND_METRIC_INTERVAL));
    }
}

void
DBImpl::WaitMergeFileFinish() {
    //    LOG_ENGINE_DEBUG_ << "Begin WaitMergeFileFinish";
    std::lock_guard<std::mutex> lck(merge_result_mutex_);
    for (auto& iter : merge_thread_results_) {
        iter.wait();
    }
    //    LOG_ENGINE_DEBUG_ << "End WaitMergeFileFinish";
}

void
DBImpl::SuspendIfFirst() {
    std::lock_guard<std::mutex> lock(suspend_build_mutex_);
    if (++live_search_num_ == 1) {
        LOG_ENGINE_TRACE_ << "live_search_num_: " << live_search_num_;
        knowhere::BuilderSuspend();
    }
}

void
DBImpl::ResumeIfLast() {
    std::lock_guard<std::mutex> lock(suspend_build_mutex_);
    if (--live_search_num_ == 0) {
        LOG_ENGINE_TRACE_ << "live_search_num_: " << live_search_num_;
        knowhere::BuildResume();
    }
}

void
DBImpl::IncreaseLiveBuildTaskNum() {
    std::lock_guard<std::mutex> lock(live_build_count_mutex_);
    ++live_build_num_;
    server_status.indexing = live_build_num_ > 0;
}

void
DBImpl::DecreaseLiveBuildTaskNum() {
    std::lock_guard<std::mutex> lock(live_build_count_mutex_);
    --live_build_num_;
    server_status.indexing = live_build_num_ > 0;
}

void
DBImpl::ConfigUpdate(const std::string& name) {
    if (name == "storage.auto_flush_interval") {
        options_.auto_flush_interval_ = config.storage.auto_flush_interval();
    }
}

}  // namespace engine
}  // namespace milvus
